import torch
import torch.utils.model_zoo as model_zoo
import os
import logging
from collections import OrderedDict


def load_checkpoint(model, checkpoint_path, use_ema=False):
    if checkpoint_path and os.path.isfile(checkpoint_path):
        checkpoint = torch.load(checkpoint_path)
        state_dict_key = ''
        if isinstance(checkpoint, dict):
            state_dict_key = 'state_dict'
            if use_ema and 'state_dict_ema' in checkpoint:
                state_dict_key = 'state_dict_ema'
        if state_dict_key and state_dict_key in checkpoint:
            new_state_dict = OrderedDict()
            for k, v in checkpoint[state_dict_key].items():
                # strip `module.` prefix
                name = k[7:] if k.startswith('module') else k
                new_state_dict[name] = v
            model.load_state_dict(new_state_dict)
        else:
            model.load_state_dict(checkpoint)
        logging.info("Loaded {} from checkpoint '{}'".format(state_dict_key or 'weights', checkpoint_path))
    else:
        logging.error("No checkpoint found at '{}'".format(checkpoint_path))
        raise FileNotFoundError()


# def resume_checkpoint(model, checkpoint_path):
#     optimizer_state = None
#     resume_epoch = None
#     resume_stage = None
#     if isinstance(checkpoint_path, list):
#         for idx, checkpoint in enumerate(checkpoint_path):
#             _process_checkpoint(model, checkpoint_path, stage=idx)
#     else:
#         _process_checkpoint(model, checkpoint_path)
#
#
# def _process_checkpoint(model, checkpoint_path, stage=None):
#     if os.path.isfile(checkpoint_path):
#         checkpoint = torch.load(checkpoint_path, map_location='cpu')
#         if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
#             new_state_dict = OrderedDict()
#             for k, v in checkpoint['state_dict'].items():
#                 name = k[7:] if k.startswith('module') else k
#                 new_state_dict[name] = v
#             model_dict = model.module.state_dict() if hasattr(model,
#                                                               'module') else model.state_dict()
#             if stage is not None:
#                 if stage == 0:
#                     stem_dict = {k: v for k, v in new_state_dict.items() if (k.startswith('_stem'))}
#                 new_state_dict = {k: v for k, v in new_state_dict.items() if
#                                   (k.startswith('_blocks.{}.'.format(stage)))}
#                 if stage == 0:
#                     new_state_dict.update(stem_dict)
#
#
#             model_dict.update(new_state_dict)
#             # print('loaded keys:', [k[:10] for k in new_state_dict])
#
#             model.load_state_dict(model_dict)
#             if 'optimizer' in checkpoint:
#                 optimizer_state = checkpoint['optimizer']
#             if 'stage' in checkpoint:
#                 resume_stage = checkpoint['stage']
#             if 'epoch' in checkpoint:
#                 resume_epoch = checkpoint['epoch']
#                 if 'version' in checkpoint and checkpoint['version'] > 1:
#                     resume_epoch += 1  # start at the next epoch, old checkpoints incremented before save
#             logging.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, checkpoint['epoch']))
#         else:
#             model.load_state_dict(checkpoint)
#             logging.info("Loaded checkpoint '{}'".format(checkpoint_path))
#         return optimizer_state, resume_epoch, resume_stage
#     else:
#         logging.error("No checkpoint found at '{}'".format(checkpoint_path))
#         raise FileNotFoundError()

def resume_checkpoint(model, checkpoint_path):
    optimizer_state = None
    resume_epoch = None
    resume_stage = None
    if os.path.isfile(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location='cpu')
        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
            new_state_dict = OrderedDict()
            for k, v in checkpoint['state_dict'].items():
                name = k[7:] if k.startswith('module') else k
                new_state_dict[name] = v
            model_dict = model.module.state_dict() if hasattr(model,
                                                              'module') else model.state_dict()
            # new_state_dict = {k: v for k, v in new_state_dict.items() if
            #                   (k in model_dict and not k.startswith('_blocks.2.'))}
            model_dict.update(new_state_dict)
            # print('loaded keys:', [k for k in new_state_dict])

            model.load_state_dict(model_dict)
            if 'optimizer' in checkpoint:
                optimizer_state = checkpoint['optimizer']
            if 'stage' in checkpoint:
                resume_stage = checkpoint['stage']
            if 'epoch' in checkpoint:
                resume_epoch = checkpoint['epoch']
                if 'version' in checkpoint and checkpoint['version'] > 1:
                    resume_epoch += 1  # start at the next epoch, old checkpoints incremented before save
            logging.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, checkpoint['epoch']))
        else:
            model.load_state_dict(checkpoint)
            logging.info("Loaded checkpoint '{}'".format(checkpoint_path))
        return optimizer_state, resume_epoch, resume_stage
    else:
        logging.error("No checkpoint found at '{}'".format(checkpoint_path))
        raise FileNotFoundError()


def load_pretrained(model, default_cfg, num_classes=1000, in_chans=3, filter_fn=None):
    if 'url' not in default_cfg or not default_cfg['url']:
        logging.warning("Pretrained model URL is invalid, using random initialization.")
        return

    state_dict = model_zoo.load_url(default_cfg['url'])

    if in_chans == 1:
        conv1_name = default_cfg['first_conv']
        logging.info('Converting first conv (%s) from 3 to 1 channel' % conv1_name)
        conv1_weight = state_dict[conv1_name + '.weight']
        state_dict[conv1_name + '.weight'] = conv1_weight.sum(dim=1, keepdim=True)
    elif in_chans != 3:
        assert False, "Invalid in_chans for pretrained weights"

    strict = True
    classifier_name = default_cfg['classifier']
    if num_classes == 1000 and default_cfg['num_classes'] == 1001:
        # special case for imagenet trained models with extra background class in pretrained weights
        classifier_weight = state_dict[classifier_name + '.weight']
        state_dict[classifier_name + '.weight'] = classifier_weight[1:]
        classifier_bias = state_dict[classifier_name + '.bias']
        state_dict[classifier_name + '.bias'] = classifier_bias[1:]
    elif num_classes != default_cfg['num_classes']:
        # completely discard fully connected for all other differences between pretrained and created model
        del state_dict[classifier_name + '.weight']
        del state_dict[classifier_name + '.bias']
        strict = False

    if filter_fn is not None:
        state_dict = filter_fn(state_dict)

    model.load_state_dict(state_dict, strict=strict)






